from collections import Counter
import pandas as pd
import json
from sklearn.model_selection import train_test_split
import os
from __00__config import Config


def preprocess_data(datapath, train_path, test_path, dev_path, json_path):
	"""
	预处理数据：过滤标签、数值化、划分数据集

	Args:
		datapath: 原始数据路径
		train_path: 训练集保存路径
		test_path: 测试集保存路径
		dev_path: 验证集保存路径
		json_path: 标签映射文件保存路径
	"""
	# 1. 读取数据
	df = pd.read_json(datapath, encoding='utf-8', lines=True)
	print("原始数据信息:")
	print(df.info())

	# 2. 只保留question和label列，并重命名列
	df = df[['question', 'label']].rename(columns={'question': 'questions', 'label': 'labels'})

	# 3. 查看原始标签分布
	original_counter = Counter(df['labels'])
	print("\n原始标签分布:")
	for label, count in original_counter.most_common():
		print(f"{label}: {count}")

	# 4. 舍弃急诊科
	# 急诊科样本数量太少了，加入训练会因为样本不均衡导致训练效果很差
	df = df[df['labels'] != '急诊科']
	print(f"\n舍弃急诊科后剩余数据量: {len(df)}")

	# 5. 查看过滤后的标签分布
	filtered_counter = Counter(df['labels'])
	print("\n过滤后标签分布:")
	for label, count in filtered_counter.most_common():
		print(f"{label}: {count}")

	# 6. 创建标签映射字典（数值化）
	unique_labels = sorted(df['labels'].unique())
	label_to_idx = {label: idx for idx, label in enumerate(unique_labels)}
	idx_to_label = {idx: label for idx, label in enumerate(unique_labels)}

	print(f"\n标签映射关系:")
	for label, idx in label_to_idx.items():
		print(f"{label} -> {idx}")

	# 7. 将标签数值化
	df['labels_idx'] = df['labels'].map(label_to_idx)

	# 8. 按标签等比例划分数据集 (6:2:2)
	# 先划分训练集和临时集 (6:4)
	train_df, temp_df = train_test_split(
		df,
		test_size=0.2,
		random_state=42,
		stratify=df['labels_idx']  # 等比例分层抽样
	)

	# 再将临时集划分为验证集和测试集 (2:2)
	dev_df, test_df = train_test_split(
		temp_df,
		test_size=0.5,
		random_state=42,
		stratify=temp_df['labels_idx']  # 等比例分层抽样
	)

	# 9. 只保留需要的列并重命名（使用数值化标签）
	train_df = train_df[['questions', 'labels_idx']].rename(columns={'labels_idx': 'labels'})
	dev_df = dev_df[['questions', 'labels_idx']].rename(columns={'labels_idx': 'labels'})
	test_df = test_df[['questions', 'labels_idx']].rename(columns={'labels_idx': 'labels'})

	# 10. 保存为CSV格式
	train_df.to_csv(train_path, index=False, encoding='utf-8')
	dev_df.to_csv(dev_path, index=False, encoding='utf-8')
	test_df.to_csv(test_path, index=False, encoding='utf-8')

	# 11. 保存标签映射文件
	label_mapping = {
		'label_to_idx': label_to_idx,
		'idx_to_label': idx_to_label
	}

	with open(json_path, 'w', encoding='utf-8') as f:
		json.dump(label_mapping, f, ensure_ascii=False, indent=2)

	# 12. 打印统计信息
	print(f"\n数据集划分完成:")
	print(f"训练集: {len(train_df)} 条 ({len(train_df) / len(df) * 100:.1f}%)")
	print(f"验证集: {len(dev_df)} 条 ({len(dev_df) / len(df) * 100:.1f}%)")
	print(f"测试集: {len(test_df)} 条 ({len(test_df) / len(df) * 100:.1f}%)")

	print(f"\n训练集标签分布:")
	train_counter = Counter(train_df['labels'])
	for idx, count in train_counter.most_common():
		label_name = idx_to_label[idx]
		print(f"{label_name}({idx}): {count}")

	print(f"\n文件保存路径:")
	print(f"训练集: {train_path}")
	print(f"验证集: {dev_path}")
	print(f"测试集: {test_path}")
	print(f"标签映射: {json_path}")

	# 13. 显示CSV文件样例
	print(f"\n训练集前3条数据样例:")
	print(train_df.head(3))


def check_data(data_path):
	df = pd.read_csv(data_path, encoding='utf-8', sep=',')
	# 查看数据信息
	# print(df.info())
	# 查看标签比例
	counter = Counter(df['labels'])
	# print(counter)
	for label, count in counter.most_common():
		print(f'{label} -> {(count / len(df["labels"])) * 100:.2f}%')


if __name__ == '__main__':
	# 假设您的Config类中有这些路径配置
	config = Config()

	# # 调用预处理函数
	# preprocess_data(
	# 	datapath=config.source_data,
	# 	train_path=config.processed_train_data,  # 需要在Config中定义，如 "data/train.csv"
	# 	test_path=config.processed_test_data,  # 需要在Config中定义，如 "data/test.csv"
	# 	dev_path=config.processed_dev_data,  # 需要在Config中定义，如 "data/dev.csv"
	# 	json_path=config.label_mapping  # 需要在Config中定义，如 "data/label_mapping.json"
	# )
	check_data(config.processed_train_data)
	# check_data(config.processed_test_data)
	# check_data(config.processed_dev_data)
